module Main where

import Prelude ()
import StdImports
import System.Environment
import System.CPUTime
import System.Timeout
import Text.Printf

import Zeno.Core
import Zeno.HaskellParser
import Zeno.FL  
import Zeno.Solver
import Zeno.Debug
import Zeno.ProofLogic

conDepthDefault = 3

arg :: String -> IO [String]
arg name = do
  as <- getArgs
  return $ map (\i -> as !! (i + 1)) (elemIndices name as) 

(<++>) :: (Monoid a, Monad m) => m a -> m a -> m a
x <++> y = (liftM mappend x) `ap` y

main :: IO ()
main = do
  args <- getArgs
  if args == ["!"]
    then test
    else if args == ["comp"]
      then comparison
    else if length args == 1
      then zeno Nothing Nothing "Z.hs" (args !! 0)
    else do
      filepath <- arg "-f"
      [goal] <- arg "-g"
      ind_steps <- arg "-i"
      limit <- arg "-t"
      con_steps <- arg "-c"
      
      let filepath' = if null filepath 
            then "Z.hs" 
            else head filepath
      
          ind = if null ind_steps 
            then Nothing 
            else Just (read (head ind_steps))
            
          con = if null con_steps
            then Nothing
            else Just (read (head con_steps))
              
      let computation = zeno ind con filepath' goal
      
      case limit of
        [] -> computation
        [limit_s] -> do
          -- Hard-code 2 secs parsing time
          let limit = ((read limit_s) * (10^6)) + 2 
          maybe_finished <- timeout limit computation
          case maybe_finished of
            Just _ -> return ()
            Nothing ->
              putStrLn $ "Computation timed out (" ++ limit_s ++ "s)."

z :: String -> IO ()
z = zeno Nothing Nothing "Z.hs"

loadZeno :: String -> IO Zeno
loadZeno filename = do
  hsfile <- readFile filename
  create_zeno <- do
    hs <- parseHaskellFile filename
    return $ hs >> findAndParseFL hsfile
  return $ execState create_zeno emptyZeno
  
comparison :: IO () 
comparison = do
  zeno <- loadZeno "Comparison.hs"
  
  forM_ (zeno `seq` [1..75]) $ \i -> do
    start <- getCPUTime
    
    let proofs = flip map [1..] $ \ind -> do
          random <- newStdGen
          return $ runReader 
            (solve (ind, conDepthDefault) ("p" ++ show i) random) zeno
          
    proof <- flip foldr1 proofs $ \proof rest -> do
      proof' <- proof
      case proof' of
        Unknown -> rest
        _ -> return proof'
        
    end <- getCPUTime
    let diff = (fromIntegral (end - start)) / (10^12)
    printf "%0.3f\n" (diff :: Double)
  
test :: IO ()
test = do
  let lemmas = [ ("leqadd",     3), 
                 ("dropconcat", 3),
                 ("bltake",     3),
                 ("minminmin",  3),
                 ("minusadd",   3),
                 ("droptake",   5),
                 ("nullnil",    1),
                 ("eqtrn",      3),
                 ("muladd",     4),
                 ("filterrev",  3),
                 ("takelenrev", 2),
                 ("leqfail",    3),
                 ("minusfail",  3),
                 ("minmaxfail", 3)]
                
  zeno <- loadZeno "Z.hs"
  writeIORef debugZeno $$! zeno
  start <- getCPUTime
  flip mapM_ lemmas $ \(lemma, i) -> do
    random <- newStdGen
    case runReader (solve (i, conDepthDefault) lemma random) zeno of
      Unknown -> putStrLn $ "Proof failed for " ++ lemma
      Proven {} -> putStrLn $ "Proved " ++ lemma
  end <- getCPUTime
  let diff = (fromIntegral (end - start)) / (10^12)
  printf "\nTotal computation time: %0.3f sec\n" (diff :: Double)
  
zeno :: Maybe Int ->  Maybe Int -> String -> String -> IO ()
zeno maybe_ind maybe_con filename lemma = do
  zeno <- loadZeno filename
  --print zeno
  
  writeIORef debugZeno $$! zeno
  let con = fromMaybe conDepthDefault maybe_con

  start <- getCPUTime
 
  proof <-
    case maybe_ind of
      Nothing -> do
        let proofs = flip map [1..] $ \ind -> do
              random <- newStdGen
              putStrLn $ "Trying depth " ++ show ind ++ "\n"
              return $ runReader (solve (ind, con) lemma random) zeno
          
        flip foldr1 proofs $ \proof rest -> do
          proof' <- proof
          case proof' of
            Unknown -> rest
            _ -> return proof'
            
      Just ind -> do
        random <- newStdGen
        return $ runReader (solve (ind, con) lemma random) zeno

  end <- getCPUTime
        
  let proof_s = runReader (showZProof proof) zeno
          
  putStrLn $ "\n" ++ proof_s
  putStrLn ""
  
  let diff = (fromIntegral (end - start)) / (10^12)
  printf "\nComputation time: %0.3f sec\n" (diff :: Double)
